package handlers

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/gin-gonic/gin"

	"agentdesign/models"
)

var (
	errPromptContextMissing   = errors.New("versionId or both projectId and pageId are required")
	errPromptPageHasNoVersion = errors.New("target page has no current version")
)

// HandlePrompt streams LLM responses back to the client using server-sent events.
func (h *Handler) HandlePrompt(c *gin.Context) {
	if h.prompt == nil {
		c.JSON(http.StatusServiceUnavailable, gin.H{"error": "prompt service is not configured"})
		return
	}

	var req models.PromptCompletionRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request payload"})
		return
	}

	messages, err := models.NormalizePromptMessages(req)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	versionID, err := h.resolvePromptVersion(req)
	if err != nil {
		status := http.StatusBadRequest
		if errors.Is(err, sql.ErrNoRows) {
			status = http.StatusNotFound
		} else if errors.Is(err, errPromptPageHasNoVersion) || errors.Is(err, errPromptContextMissing) {
			status = http.StatusBadRequest
		} else {
			status = http.StatusInternalServerError
		}
		c.JSON(status, gin.H{"error": err.Error()})
		return
	}

	userContent := strings.TrimSpace(latestUserMessage(messages))
	if userContent == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "user prompt is empty"})
		return
	}

	req.VersionID = versionID

	project := h.findProject(req.ProjectID)
	projectBrief := formatProjectBrief(project)

	promptRequest, err := h.service.CreatePromptRequest(models.PromptRequest{
		ProjectID:          optionalStringPointer(req.ProjectID),
		PageID:             optionalStringPointer(req.PageID),
		PromptText:         userContent,
		Status:             "pending",
		ModelName:          req.Model,
		GeneratedVersionID: optionalStringPointer(versionID),
	})
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"error": "failed to create prompt request",
			"msg":   err.Error(),
		})
		return
	}

	promptRequestID := promptRequest.ID

	if err := h.persistUserPrompt(versionID, userContent, &promptRequestID); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"error": "failed to store prompt message",
			"msg":   err.Error(),
		})
		return
	}

	preparedMessages, err := h.preparePromptMessages(req, messages, project, projectBrief, userContent)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"error": "failed to prepare prompt messages",
			"msg":   err.Error(),
		})
		return
	}

	flusher, ok := c.Writer.(http.Flusher)
	if !ok {
		c.JSON(http.StatusInternalServerError, gin.H{
			"error": "streaming not supported by server",
			"msg":   "streaming not supported by server",
		})
		return
	}

	headers := c.Writer.Header()
	headers.Set("Content-Type", "text/event-stream")
	headers.Set("Cache-Control", "no-cache")
	headers.Set("Connection", "keep-alive")

	c.Status(http.StatusOK)
	flusher.Flush()

	var (
		assistantBuffer strings.Builder
		assistantStored bool
		promptCompleted bool
	)

	send := func(chunk models.PromptStreamChunk) error {
		if chunk.Type == models.PromptStreamChunkTypeToken {
			assistantBuffer.WriteString(chunk.Data)
		}

		if chunk.Type == models.PromptStreamChunkTypeDone && !assistantStored {
			content := strings.TrimSpace(assistantBuffer.String())
			if err := h.persistAssistantResponse(versionID, content, &promptRequestID); err != nil {
				return err
			}
			lang, snippet, err := h.captureCodeSnippets(versionID, content)
			if err != nil {
				return err
			}

			if promptRequestID != "" {
				responsePayload := map[string]any{
					"content": content,
				}
				if snippet != "" {
					responsePayload["snippet"] = map[string]string{
						"language": lang,
						"content":  snippet,
					}
				}
				if data, err := json.Marshal(responsePayload); err == nil {
					_ = h.service.MarkPromptRequestCompleted(promptRequestID, data)
					promptCompleted = true
				}
			}

			assistantStored = true
		}

		payload, err := json.Marshal(chunk)
		if err != nil {
			return err
		}
		if _, err := c.Writer.Write([]byte("data: ")); err != nil {
			return err
		}
		if _, err := c.Writer.Write(payload); err != nil {
			return err
		}
		if _, err := c.Writer.Write([]byte("\n\n")); err != nil {
			return err
		}
		flusher.Flush()
		return nil
	}

	err = h.prompt.StreamChatCompletion(
		c.Request.Context(),
		models.PromptCompletionInput{
			ProjectID:   req.ProjectID,
			PageID:      req.PageID,
			VersionID:   versionID,
			Model:       req.Model,
			Messages:    preparedMessages,
			Temperature: req.Temperature,
			MaxTokens:   req.MaxTokens,
			TopP:        req.TopP,
		},
		send,
	)
	if err != nil {
		_ = send(models.PromptStreamChunk{
			Type:  models.PromptStreamChunkTypeError,
			Error: err.Error(),
		})
		if !promptCompleted && promptRequestID != "" {
			_ = h.service.MarkPromptRequestFailed(promptRequestID, err.Error())
		}
		return
	}
}

func (h *Handler) resolvePromptVersion(req models.PromptCompletionRequest) (string, error) {
	if v := strings.TrimSpace(req.VersionID); v != "" {
		return v, nil
	}

	projectID := strings.TrimSpace(req.ProjectID)
	pageID := strings.TrimSpace(req.PageID)
	if projectID == "" || pageID == "" {
		return "", errPromptContextMissing
	}

	page, err := h.service.GetProjectPage(projectID, pageID)
	if err != nil {
		return "", err
	}

	if strings.TrimSpace(page.CurrentVersionID) == "" {
		return "", errPromptPageHasNoVersion
	}

	return page.CurrentVersionID, nil
}

func (h *Handler) persistUserPrompt(versionID, content string, promptRequestID *string) error {
	if content == "" {
		return nil
	}

	_, err := h.service.CreateCanvasMessage(models.CanvasMessage{
		VersionID:       versionID,
		Author:          "user",
		Content:         content,
		PromptRequestID: promptRequestID,
	})
	return err
}

func (h *Handler) persistAssistantResponse(versionID, content string, promptRequestID *string) error {
	content = strings.TrimSpace(content)
	if content == "" {
		return nil
	}

	_, err := h.service.CreateCanvasMessage(models.CanvasMessage{
		VersionID:       versionID,
		Author:          "assistant",
		Content:         content,
		PromptRequestID: promptRequestID,
	})
	return err
}

func latestUserMessage(messages []models.PromptMessage) string {
	for i := len(messages) - 1; i >= 0; i-- {
		if strings.ToLower(strings.TrimSpace(messages[i].Role)) == "user" {
			return messages[i].Content
		}
	}
	return ""
}

var codeFenceRegex = regexp.MustCompile("(?s)```([\\w.+-]*)\\s*\\n(.*?)```")

func (h *Handler) captureCodeSnippets(versionID, content string) (string, string, error) {
	snippets := extractCodeSnippets(content)
	if len(snippets) == 0 {
		return "", "", nil
	}

	language, merged := consolidateSnippets(snippets)
	if merged == "" {
		return "", "", nil
	}

	payload := map[string]string{
		"type":     "code",
		"language": language,
		"content":  merged,
	}
	componentJSON, err := json.Marshal(payload)
	if err != nil {
		return "", "", err
	}

	label := "Code Snippet"
	if language != "" {
		label = fmt.Sprintf("%s Snippet", strings.ToUpper(language))
	}

	_, err = h.service.CreatePageSection(models.PageSection{
		VersionID:     versionID,
		SectionKey:    fmt.Sprintf("code-snippet-%d", time.Now().UnixMilli()),
		Label:         label,
		Description:   "Generated code block from assistant response",
		Position:      1000,
		ComponentJSON: componentJSON,
	})
	if err != nil {
		return "", "", err
	}

	return language, merged, nil
}

type codeSnippet struct {
	Language string
	Content  string
}

func extractCodeSnippets(content string) []codeSnippet {
	matches := codeFenceRegex.FindAllStringSubmatch(content, -1)
	snippets := make([]codeSnippet, 0, len(matches))

	for _, match := range matches {
		if len(match) < 3 {
			continue
		}
		language := strings.TrimSpace(match[1])
		body := strings.TrimSpace(match[2])
		if body == "" {
			continue
		}
		snippets = append(snippets, codeSnippet{
			Language: language,
			Content:  body,
		})
	}

	return snippets
}

func consolidateSnippets(snippets []codeSnippet) (string, string) {
	if len(snippets) == 0 {
		return "", ""
	}

	var builder strings.Builder
	language := ""

	for idx, snippet := range snippets {
		if language == "" && snippet.Language != "" {
			language = snippet.Language
		}

		if idx > 0 {
			builder.WriteString("\n\n")
			builder.WriteString("<!-- snippet separator -->\n\n")
		}

		builder.WriteString(snippet.Content)
	}

	if language == "" {
		language = "markdown"
	}

	return language, strings.TrimSpace(builder.String())
}

func (h *Handler) preparePromptMessages(
	req models.PromptCompletionRequest,
	messages []models.PromptMessage,
	project *models.Project,
	projectBrief string,
	userContent string,
) ([]models.PromptMessage, error) {
	stage := determinePromptStage(messages)
	instruction, err := h.service.GetPromptInstructionBySlug(templateSlugForStage(stage))
	if err != nil {
		return nil, err
	}

	if strings.TrimSpace(projectBrief) == "" {
		projectBrief = "Not provided."
	}

	values := map[string]string{
		"REQUEST":          userContent,
		"PROJECT_ID":       req.ProjectID,
		"PAGE_ID":          req.PageID,
		"VERSION_ID":       req.VersionID,
		"PROJECT_BRIEF":    projectBrief,
		"PROJECT_NAME":     "",
		"PROJECT_STATUS":   "",
		"PROJECT_PLATFORM": "",
	}

	if project != nil {
		values["PROJECT_NAME"] = project.Name
		values["PROJECT_STATUS"] = project.Status
		values["PROJECT_PLATFORM"] = project.Platform
	}

	templatedContent := injectTemplateValues(instruction.TemplateText, values)

	prepared := []models.PromptMessage{
		{
			Role:    "system",
			Content: instruction.Instruction,
		},
	}

	lastUserIdx := lastUserMessageIndex(messages)
	for idx, msg := range messages {
		if idx == lastUserIdx {
			continue
		}
		prepared = append(prepared, msg)
	}

	prepared = append(prepared, models.PromptMessage{
		Role:    "user",
		Content: templatedContent,
	})

	return prepared, nil
}

type promptStage string

const (
	stageInitial   promptStage = "initial"
	stageIterative promptStage = "iterative"
)

func determinePromptStage(messages []models.PromptMessage) promptStage {
	assistantSeen := false
	userCount := 0

	for _, msg := range messages {
		role := strings.ToLower(strings.TrimSpace(msg.Role))
		switch role {
		case "assistant":
			assistantSeen = true
		case "user":
			userCount++
		}
	}

	if assistantSeen && userCount > 1 {
		return stageIterative
	}

	return stageInitial
}

func templateSlugForStage(stage promptStage) string {
	switch stage {
	case stageIterative:
		return "component-refiner"
	default:
		return "ux-design-default"
	}
}

func injectTemplateValues(template string, values map[string]string) string {
	result := template
	for key, value := range values {
		placeholder := "{{" + strings.ToUpper(key) + "}}"
		result = strings.ReplaceAll(result, placeholder, value)
	}
	return result
}

func lastUserMessageIndex(messages []models.PromptMessage) int {
	for i := len(messages) - 1; i >= 0; i-- {
		if strings.ToLower(strings.TrimSpace(messages[i].Role)) == "user" {
			return i
		}
	}
	return -1
}

func (h *Handler) findProject(projectID string) *models.Project {
	if strings.TrimSpace(projectID) == "" {
		return nil
	}

	project, err := h.service.GetProject(projectID)
	if err != nil {
		return nil
	}

	return &project
}

func formatProjectBrief(project *models.Project) string {
	if project == nil {
		return ""
	}

	var builder strings.Builder

	if project.Name != "" {
		builder.WriteString("Project: ")
		builder.WriteString(project.Name)
	}

	if project.Status != "" {
		if builder.Len() > 0 {
			builder.WriteString(" | ")
		}
		builder.WriteString("Status: ")
		builder.WriteString(project.Status)
	}

	if project.Platform != "" {
		if builder.Len() > 0 {
			builder.WriteString(" | ")
		}
		builder.WriteString("Platform: ")
		builder.WriteString(project.Platform)
	}

	if strings.TrimSpace(project.Description) != "" {
		if builder.Len() > 0 {
			builder.WriteString("\n")
		}
		builder.WriteString("Brief: ")
		builder.WriteString(project.Description)
	}

	return builder.String()
}

func optionalStringPointer(value string) *string {
	trimmed := strings.TrimSpace(value)
	if trimmed == "" {
		return nil
	}
	v := trimmed
	return &v
}
